package com.ciei.dpagm.util;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 树形数据工具类
 *
 * @author fangnuozhou
 * @date 2020/08/20
 */
public class TreeUtils {

    /**
     * 构建树形结构数据
     *
     * @param nodes     总数据集
     * @return          树形结构数据集
     */
    public static List<TreeNode> buildTree(List<TreeNode> nodes) {
        Map<Integer, List<TreeNode>> childrenMap = nodes.stream()
                .filter(node -> node.getPid() != null && node.getPid() != 0)
                .collect(Collectors.groupingBy(TreeNode::getPid));
        nodes.forEach(node -> {
            List<TreeNode> children = childrenMap.get(node.getId());
            if (children != null) {
                Collections.sort(children);
                node.setChildren(children);
            }
        });
        return nodes.stream().filter(node -> node.getPid() == null || node.getPid() == 0).sorted().collect(Collectors.toList());
    }

    /**
     * 构建树形结构数据
     *
     * @param dataList          总数据集
     * @param idMapper          取ID的方法
     * @param pidMapper         取父ID的方法
     * @param labelMapper       取显示内容的方法
     * @param <T>               泛型
     * @return                  树形结构数据集
     */
    public static <T> List<TreeNode> buildTree(List<T> dataList, Function<? super T, Integer> idMapper,
                                           Function<? super T, Integer> pidMapper, Function<? super T, String> labelMapper) {
        return buildTree(dataList, idMapper, pidMapper, labelMapper, null);
    }

    /**
     * 构建树形结构数据
     *
     * @param dataList          总数据集
     * @param idMapper          取ID的方法
     * @param pidMapper         取父ID的方法
     * @param labelMapper       取显示内容的方法
     * @param <T>               泛型
     * @return                  树形结构数据集
     */
    public static <T> List<TreeNode> buildTree(List<T> dataList, Function<? super T, Integer> idMapper,
                                               Function<? super T, Integer> pidMapper, Function<? super T, String> labelMapper,
                                               Function<? super T, Integer> orderMapper) {
        if (dataList == null || dataList.isEmpty()) {
            return new ArrayList<>();
        }

        List<TreeNode> nodes = new ArrayList<>();
        dataList.forEach(data -> {
            TreeNode node = new TreeNode(idMapper.apply(data), pidMapper.apply(data), labelMapper.apply(data));
            if (orderMapper != null) {
                node.setOrder(orderMapper.apply(data));
            }
            nodes.add(node);
        });
        return buildTree(nodes);
    }
}
